{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "99dcd77a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "输入输出形状测试通过\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from models.discriminator import Discriminator_small, Discriminator_large\n",
    "\n",
    "# 配置参数（根据实际任务调整）\n",
    "batch_size = 2\n",
    "image_size = 32  # 或256，根据影像尺寸选择\n",
    "num_timesteps = 100  # 扩散时间步数量\n",
    "\n",
    "# 构造测试数据：T1和T2均为单通道影像\n",
    "t1 = torch.randn(batch_size, 1, image_size, image_size)  # 随机模拟T1影像\n",
    "t2_real = torch.randn(batch_size, 1, image_size, image_size)  # 随机模拟真实T2影像\n",
    "t2_fake = torch.randn(batch_size, 1, image_size, image_size)  # 随机模拟生成的T2影像\n",
    "t = torch.randint(0, num_timesteps, (batch_size,))  # 随机时间步\n",
    "\n",
    "# 初始化判别器（修改后应支持nc=2）\n",
    "disc = Discriminator_small(nc=2)  # 或Discriminator_large(nc=2)\n",
    "\n",
    "# 测试forward方法输出形状\n",
    "output_real = disc(t1, t, t2_real)\n",
    "output_fake = disc(t1, t, t2_fake)\n",
    "\n",
    "assert output_real.shape == (batch_size, 1), f\"输出形状错误，应为{(batch_size, 1)}，实际为{output_real.shape}\"\n",
    "assert output_fake.shape == (batch_size, 1), f\"输出形状错误，应为{(batch_size, 1)}，实际为{output_fake.shape}\"\n",
    "print(\"输入输出形状测试通过\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "92bf055b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "真实对判别分数: -0.0000\n",
      "生成对判别分数: -0.0000\n"
     ]
    },
    {
     "ename": "AssertionError",
     "evalue": "真实/生成对判别分数无差异",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mAssertionError\u001b[0m                            Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[2], line 9\u001b[0m\n\u001b[0;32m      6\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m生成对判别分数: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mfake_score\u001b[38;5;132;01m:\u001b[39;00m\u001b[38;5;124m.4f\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m      8\u001b[0m \u001b[38;5;66;03m# 随机参数下可能无显著差异，但至少不应完全相同\u001b[39;00m\n\u001b[1;32m----> 9\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m torch\u001b[38;5;241m.\u001b[39misclose(torch\u001b[38;5;241m.\u001b[39mtensor(real_score), torch\u001b[38;5;241m.\u001b[39mtensor(fake_score), atol\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1e-3\u001b[39m), \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m真实/生成对判别分数无差异\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m     10\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m判别逻辑初步测试通过\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "\u001b[1;31mAssertionError\u001b[0m: 真实/生成对判别分数无差异"
     ]
    }
   ],
   "source": [
    "# 计算真实对和生成对的判别分数均值\n",
    "real_score = output_real.mean().item()\n",
    "fake_score = output_fake.mean().item()\n",
    "\n",
    "print(f\"真实对判别分数: {real_score:.4f}\")\n",
    "print(f\"生成对判别分数: {fake_score:.4f}\")\n",
    "\n",
    "# 随机参数下可能无显著差异，但至少不应完全相同\n",
    "assert not torch.isclose(torch.tensor(real_score), torch.tensor(fake_score), atol=1e-3), \"真实/生成对判别分数无差异\"\n",
    "print(\"判别逻辑初步测试通过\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ddd16119",
   "metadata": {},
   "outputs": [
    {
     "ename": "AssertionError",
     "evalue": "时间步对输出无影响，时间嵌入可能失效",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mAssertionError\u001b[0m                            Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[3], line 9\u001b[0m\n\u001b[0;32m      6\u001b[0m output_t0 \u001b[38;5;241m=\u001b[39m disc(t1, t0, t2)\n\u001b[0;32m      7\u001b[0m output_t10 \u001b[38;5;241m=\u001b[39m disc(t1, t10, t2)\n\u001b[1;32m----> 9\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m torch\u001b[38;5;241m.\u001b[39mallclose(output_t0, output_t10, atol\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1e-3\u001b[39m), \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m时间步对输出无影响，时间嵌入可能失效\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m     10\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m时间嵌入测试通过\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "\u001b[1;31mAssertionError\u001b[0m: 时间步对输出无影响，时间嵌入可能失效"
     ]
    }
   ],
   "source": [
    "t1 = torch.randn(batch_size, 1, image_size, image_size)\n",
    "t2 = torch.randn(batch_size, 1, image_size, image_size)\n",
    "t0 = torch.zeros(batch_size, dtype=torch.int32)  # 时间步0\n",
    "t10 = torch.full((batch_size,), 10, dtype=torch.int32)  # 时间步10\n",
    "\n",
    "output_t0 = disc(t1, t0, t2)\n",
    "output_t10 = disc(t1, t10, t2)\n",
    "\n",
    "assert not torch.allclose(output_t0, output_t10, atol=1e-3), \"时间步对输出无影响，时间嵌入可能失效\"\n",
    "print(\"时间嵌入测试通过\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b8ac847a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "特征提取功能测试通过\n"
     ]
    }
   ],
   "source": [
    "# 修正特征提取功能检查的尺寸计算逻辑\n",
    "features = disc.get_features(t1, t, t2_real)\n",
    "\n",
    "assert len(features) == 5, f\"特征图数量错误，应为5，实际为{len(features)}\"\n",
    "\n",
    "prev_size = image_size\n",
    "for i, feat in enumerate(features):\n",
    "    # 下采样仅发生在i=2,3,4（对应conv2, conv3, conv4）\n",
    "    # i=0,1: 无下采样；i>=2: 下采样次数为(i-1)次（因为从i=2开始第一次下采样）\n",
    "    if i < 2:\n",
    "        expected_size = image_size  # 前两层尺寸不变\n",
    "    else:\n",
    "        downsample_count = i - 1  # i=2→1次，i=3→2次，i=4→3次\n",
    "        expected_size = image_size // (2** downsample_count)\n",
    "    \n",
    "    assert feat.shape[2] == expected_size, f\"特征图{i}尺寸错误，应为{expected_size}，实际为{feat.shape[2]}\"\n",
    "    prev_size = feat.shape[2]\n",
    "print(\"特征提取功能测试通过\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c9a8611f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "梯度传播测试通过\n"
     ]
    }
   ],
   "source": [
    "# 切换到训练模式\n",
    "disc.train()\n",
    "\n",
    "# 构造输入并计算损失\n",
    "t1 = torch.randn(batch_size, 1, image_size, image_size, requires_grad=True)\n",
    "t2_real = torch.randn(batch_size, 1, image_size, image_size)\n",
    "t = torch.randint(0, num_timesteps, (batch_size,))\n",
    "\n",
    "output = disc(t1, t, t2_real)\n",
    "loss = torch.nn.functional.softplus(-output).mean()  # 真实对的损失\n",
    "\n",
    "# 反向传播\n",
    "loss.backward()\n",
    "\n",
    "# 检查关键层参数是否有梯度\n",
    "assert disc.start_conv.weight.grad is not None, \"start_conv层无梯度\"\n",
    "assert disc.conv1.conv1[0].weight.grad is not None, \"conv1层无梯度\"\n",
    "assert disc.end_linear.weight.grad is not None, \"end_linear层无梯度\"\n",
    "print(\"梯度传播测试通过\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cdal_cpu",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.23"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
